import argparse

def str_to_bool(value):
    if value.lower() in ('True', 'true', '1', 't', 'y', 'yes'):
        return True
    elif value.lower() in ('False', 'false', '0', 'f', 'n', 'no'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected')

class ParamInit:
    def __init__(self):
        self.param_line_args = {}
        self._create_parser()
        self.bounds = {}
        self.hard_constrains = []
        self.num_populations = 0
        self.gene_length = 0
        self.get_genetic_input()
    

    def _create_parser(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('--model', type=str, help="model name", required=True)
        parser.add_argument('--pref_batch_lower', type=int, help="pref_batch_lower", required=False, default=20)
        parser.add_argument('--pref_batch_upper', type=int, help="pref_batch_upper", required=False, default=50)
        parser.add_argument('--decode_batch_lower', type=int, help="decode_batch_lower", required=False, default=30)
        parser.add_argument('--decode_batch_upper', type=int, help="decode_batch_upper", required=False, default=150)
        parser.add_argument('--prefill_token_delay_tolerance_lower', type=int, help="prefill_token_delay_tolerance_lower", required=False, default=500)
        parser.add_argument('--prefill_token_delay_tolerance_upper', type=int, help="prefill_token_delay_tolerance_upper", required=False, default=800)
        parser.add_argument('--decode_constrains', type=float, help="decode_constrains", required=False, default=50)
        parser.add_argument('--is_firsttoken_constrained', type=str_to_bool, help="is_firsttoken_constrained", required=False, default=False)
        parser.add_argument('--firsttoken_constrains', type=float, help="firsttoken_constrains", required=False, default=1000)
        parser.add_argument('--num_populations', type=int, help="num_populations", required=False, default=4)
        parser.add_argument('--num_iterations', type=int, help="num_iterations", required=False, default=4)
        parser.add_argument('--concurrency', type=int, help="Concurrency", required=False, default=1000)
        parser.add_argument('--request_rate', type=int, help="Requestrate", required=False, default=20)
        parser.add_argument('--is_requestrate_compute', type=str_to_bool, help="Requestrate", required=False, default=False)
        parser.add_argument('--request_rate_lower', type=float, help="request_rate_lower", required=False, default=15)
        parser.add_argument('--request_rate_upper', type=float, help="request_rate_upper", required=False, default=25)
        parser.add_argument('--using_genetic', type=str_to_bool, help="Genetic Algorithm", required=False, default=False)
        parser.add_argument('--gene_length', type=int, help="gene_length", required=False, default=10)
        parser.add_argument('--is_SLO', type=str_to_bool, help="SLO", required=False, default=False)
        parser.add_argument('--is_P90', type=str_to_bool, help="P90", required=False, default=False)
        parser.add_argument('--no_constrain', type=str_to_bool, help="no_constrain", required=False, default=False)
        parser.add_argument('--is_prefixcache', type=str_to_bool, help="no_constrain", required=False, default=False)
        parser.add_argument('--is_splitfuse', type=str_to_bool, help="no_constrain", required=False, default=False)
        parser.add_argument('--data_name', type=str, help="no_constrain", required=False, default="short")
        parser.add_argument('--output_len', type=int, help="no_constrain", required=False, default=512)
        parser.add_argument('--supportSelectBatch', type=str_to_bool, help="supportSelectBatch", required=False, default=False)
        parser.add_argument('--is_speculative', type=str_to_bool, help="is_speculative", required=False, default=False)

        args = parser.parse_args()
        self.param_line_args = vars(args)

    def get_genetic_input(self):
         # 初始化返回值
        bounds = {
                    'Prefill BatchSize': [
                        float(self.get_args("pref_batch_lower")),
                        float(self.get_args("pref_batch_upper"))
                    ],
                    'Decode BatchSize': [
                        float(self.get_args("decode_batch_lower")),
                        float(self.get_args("decode_batch_upper"))
                    ],
                    'SelectBatch Prefill Delay Tolerance': [
                        float(self.get_args("prefill_token_delay_tolerance_lower")),
                        float(self.get_args("prefill_token_delay_tolerance_upper"))
                    ],
                    'Request Rate': [
                        float(self.get_args("request_rate_lower")),
                        float(self.get_args("request_rate_upper"))
                    ]
                }
        if self.get_args("is_firsttoken_constrained"):
            self.hard_constrains = [float(self.get_args("decode_constrains")), float(self.get_args("firsttoken_constrains"))]
        else:
            self.hard_constrains = [float(self.get_args("decode_constrains"))]
        self.num_populations = self.get_args("num_populations")
        self.gene_length = self.get_args("gene_length")

        # 返回用户输入的数据
        # return bounds, hard_constrains, num_populations, gene_length
    
    def get_args(self, key):
        return self.param_line_args.get(key)

# config = ParamInit()
# print(config.param_line_args)
